import gradio as gr
import pickle
import numpy as np
from PIL import Image

# 加载保存的KNN模型
with open('best_knn_model.pkl', 'rb') as model_file:
    knn_model = pickle.load(model_file)
from PIL import ImageOps

# 定义预测函数
def predict_digit(image):
    # 如果 image 是字典，提取图像数据
    if isinstance(image, dict):
        image = image.get('composite')  # 使用 'composite' 键提取图像数据
        if image is None:
            raise ValueError("Image data not found in 'composite'")

    # 如果图像有透明度通道，取第四通道；否则处理为灰度图
    if image.shape[2] == 4:  # 检查是否有 Alpha 通道
        alpha_channel = image[:, :, 3]  # 获取第四通道，表示透明度
    else:
        # 如果没有透明度通道，取RGB平均值转换为灰度图
        alpha_channel = np.mean(image[:, :, :3], axis=2)
    
    # 将图像转换为PIL图像并调整大小为8x8
    pil_image = Image.fromarray(alpha_channel).resize((8, 8))
    
    # 对图像进行反转和二值化处理，使数字更加清晰
    pil_image = ImageOps.invert(pil_image.convert('L'))  # 反转颜色并转换为灰度图
    pil_image = pil_image.point(lambda x: 0 if x < 128 else 255, 'L')  # 二值化处理（灰度值）
    
    # 打印转换后的PIL图像信息
    print(f"Processed PIL image (after binarization): {pil_image}")

    # 转换为NumPy数组并将其展平为1D数组
    image_np = np.array(pil_image).reshape(1, -1)

    # 打印转换后的NumPy数组
    print(f"Processed NumPy image: {image_np}")

    # 归一化像素值到[0, 16]区间（digits 数据集像素值范围为 0-16）
    image_np = (image_np / 255.0) * 16

    # 使用KNN模型进行预测
    prediction = knn_model.predict(image_np)
    return int(prediction[0])

# 创建 Gradio 接口
with gr.Blocks() as demo:
    gr.Markdown("# 手写数字识别应用")

    with gr.Row():
        with gr.Column():
            # 左侧：用户绘制数字的输入
            sketchpad = gr.Sketchpad(label="Draw a digit", width=280, height=280)  # 设置画布大小
            clear_button = gr.Button("Clear")  # 增加 Clear 按钮
            submit_button = gr.Button("Submit")  # 增加 Submit 按钮
        with gr.Column():
            # 右侧：预测结果
            label_output = gr.Label(label="Prediction Output")

    # 将 "Submit" 按钮绑定到预测函数上
    submit_button.click(fn=predict_digit, inputs=sketchpad, outputs=label_output)

    # 将 "Clear" 按钮绑定到清空绘图板
    clear_button.click(lambda: None, None, sketchpad)

# 启动 Gradio 接口
demo.queue().launch()

